# -*- coding: utf-8 -*-
# FileName:     zaudio.py
# time:         23/4/23 023 下午 1:15
# Author:       Zhou Hang
# Description:  基于PyAudio, block模式播放音频
import pyaudio
from src.utils.common import *
from src.utils.zwave import *

SAMPLE_BUFFER_SIZE = 4096
BYTES_PER_SAMPLE = 2


class StdAudio:
    _instance = None

    def __new__(cls, *args, **kw):
        if cls._instance is None:
            cls._instance = object.__new__(cls, *args, **kw)
        return cls._instance

    def __init__(self, samplewidth=2, channels=1, samplerate=SAMPLE_RATE):
        """
        :param samplewidth: 采样精度, 可以为
        :param channels: 通道数
        :param samplerate: 采样率
        """
        self.save_mode = False  # 默认关闭保存模式
        self.new_record_fn = ""
        self.new_record = None  # ReadWave类型值，若打开保存模式，必不为None
        self.samplerate = samplerate
        self.samplewidth = samplewidth
        if samplewidth == 2:
            self.maxnbit = MAX16BIT
        elif samplewidth == 1:
            self.maxnbit = MAX8BIT
        self.p = pyaudio.PyAudio()
        self.stream = self.p.open(format=pyaudio.get_format_from_width(samplewidth),
                                  channels=channels,
                                  rate=samplerate,
                                  output=True,
                                  # frames_per_buffer=1024,
                                  # stream_callback=self.play_callback
                                  )
        self.buffer = bytes()
        self.buffer_capacity = SAMPLE_BUFFER_SIZE * BYTES_PER_SAMPLE
        self.buffer_size = 0

    def set_save_mode(self, switch, filename: str = ""):
        """
        调用这个函数后，play_float播放的声音会被保存到wav文件中
        :param switch: True-开启保存模式，False-关闭保存模式
        :param filename: 如果设置了这个值，播放的声音会被保存到这个文件中，否则生成在默认位置
        """
        self.save_mode = switch
        if switch:
            print("开始录音")
            if filename == "":
                self.new_record_fn = f"{DEFAULT_RECORD_DIR}/{gen_unique_name()}.wav"
            else:
                self.new_record_fn = filename
            self.new_record = WriteWave(self.new_record_fn)
        else:
            print("结束录音")
            if self.new_record is not None:
                self.__finish_record()
                self.new_record = None
                self.new_record_fn = ""

    def __finish_record(self):
        self.new_record.close()

    def play_float(self, sample: float):
        """
        播放一个float类型数据表示的音频，本质上是将float映射到Short范围上，实际播放的还是一个Short精度的值
        :param sample:
        :return:
        """
        sample = float2sample(sample)
        self.buffer += sample
        self.buffer_size += self.samplewidth
        if self.buffer_size >= self.buffer_capacity:  # 缓冲区写满后才会向pystream中写入信息
            self.stream.write(self.buffer)
            if self.save_mode:  # 保存模式，在这里进行保存
                assert self.new_record
                self.new_record.write_frames(self.buffer)
            self.buffer_size = 0
            self.buffer = bytes()

    def close(self):
        self.stream.stop_stream()
        self.stream.close()
        self.p.terminate()

    @staticmethod
    def play(filename):
        """
        播放指定文件的音频
        :param filename: 指定文件
        :return:
        """
        wav = ReadWave(filename)
        local_p = pyaudio.PyAudio()
        local_stream = local_p.open(format=local_p.get_format_from_width(wav.sampwidth),
                                    channels=wav.nchannels,
                                    rate=wav.framerate,
                                    output=True
                                    )
        local_stream.write(wav.raw_data)  # TODO: 这里会不会出现raw_data太大导致没有播放完全？
        local_stream.stop_stream()
        local_stream.close()
        local_p.terminate()


class StdRecord:
    def __init__(self, samplewidth=2, channels=1, samplerate=44100) -> None:
        self.audio = pyaudio.PyAudio()
        self.samplerate = samplerate
        # print(get_usb_microphone_index(self.p))
        self.stream = self.audio.open(format=self.audio.get_format_from_width(samplewidth),
                                      channels=channels,
                                      rate=samplerate,
                                      input=True,
                                      input_device_index=self.get_usb_microphone_index(),
                                      frames_per_buffer=CHUNK
                                      )

    def get_usb_microphone_index(self):
        """
        获得usb麦克风设备号, 若找不到则使用默认设备
        """
        device_index = 0
        for i in range(self.audio.get_device_count()):
            device_info = self.audio.get_device_info_by_index(i)
            # print(device_info)
            if device_info.get('name', '').find('USB') != -1:
                device_index = device_info.get('index')
        return device_index

    def record(self, filename: str, seconds: int):
        ww = WriteWave(filename)
        for i in range(0, int(self.samplerate * seconds / CHUNK)):
            data = self.stream.read(CHUNK)
            ww.write_frames(data)
        ww.close()

    def get_sound_data(self, f, seconds=None):
        """
        对收到的声音做处理
        :param f: 回调函数，对收到的声音数据处理
        :param seconds: 若为None, 则一直收音并处理
        """
        if seconds is None:
            while True:
                data = self.stream.read(CHUNK, exception_on_overflow=False)
                f(data)
        else:
            for i in range(0, int(self.samplerate * seconds / CHUNK)):
                data = self.stream.read(CHUNK, exception_on_overflow=False)
                f(data)

    def get_sound_data_wait(self):
        """将获得的数据yield出去的版本"""
        while True:
            data = self.stream.read(CHUNK, exception_on_overflow=False)
            yield data

    def close(self):
        self.stream.stop_stream()
        self.stream.close()
        self.audio.terminate()


def main():
    # 测试StdAudio是不是单例
    s1 = StdAudio()
    s2 = StdAudio()
    assert s1 == s2


if __name__ == "__main__":
    main()
